package com.dx.chat.kafka.partition;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.kafka.clients.producer.Partitioner;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.utils.Utils;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;

/**
 * 自定义分区策略
 *
 * @author xufeng
 * @since 21.06.2023
 */
@Component
public class CustomizePartitioner implements Partitioner {

    private static final Log logger = LogFactory.getLog(CustomizePartitioner.class);

    /**
     * 自定义分区策略
     */
    @Override
    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
        // 获取topic的分区列表
        List<PartitionInfo> partitionInfoList = cluster.partitionsForTopic(topic);
        if (keyBytes == null) {
            // 默认为0
            return 0;
        }
        // 取key的绝对值对总partition求余
        return Utils.toPositive(Utils.murmur2(keyBytes)) % partitionInfoList.size();
    }

    /**
     * 在分区程序关闭时调用
     */
    @Override
    public void close() {
        logger.info("partitioner close ...");
    }

    /**
     * 做必要的初始化工作
     */
    @Override
    public void configure(Map<String, ?> configs) {
        logger.info("partitioner init ...");
    }
}
